""" 数据库连接中间件
@Project: Rebar
@File: exception_middleware.py
@Author: hzc
@Date: 2023/3/21 16:56
"""
import traceback

from starlette.requests import Request

from model.db.connect import tenant_session_scope
from model.db.mapper.log_mapper import add_log
from util import error_res, ErrorCode
from model.db import connect


class DatabaseHandlerMiddleware:
    def __init__(self, app=None):
        self.app = app

    async def __call__(self, request: Request, call_next):
        """
        :param request:
        :param call_next:
        :return:
        """
        project_id = request.headers.get("projectID")
        if project_id is not None:
            try:
                # 连接项目数据库
                connect.projSession = tenant_session_scope("project" + str(project_id))
            except Exception as e:
                return error_res(ErrorCode.SERVER, '项目数据库无法连接')
        else:
            # 连接公共数据库
            connect.projSession = tenant_session_scope("reservoir")

        # 让路由处理请求并获取响应
        response = await call_next(request)
        # 在这里可以对响应进行拦截和处理
        # 关闭数据库连接
        connect.projSession.close()
        return response
